// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <krpc/rpc/restful_service.h>
#include <kllm/core/km_context.h>

namespace kllm {

    struct HealthProcessor : public krpc::RestfulProcessor {
        HealthProcessor() = default;
        ~HealthProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;
    };

    struct MetricProcessor : public krpc::RestfulProcessor {

        MetricProcessor(KMContext *c) : context(c) {}

        ~MetricProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct PropsProcessor : public krpc::RestfulProcessor {

        PropsProcessor(KMContext *c) : context(c) {}

        ~PropsProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    private:
        void get(const krpc::RestfulRequest *request, krpc::RestfulResponse *response);

        void post(const krpc::RestfulRequest *request, krpc::RestfulResponse *response);
    };

    struct ModelsProcessor : public krpc::RestfulProcessor {

        ModelsProcessor(KMContext *c) : context(c) {}

        ~ModelsProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct CompletionsProcessor : public krpc::RestfulProcessor {

        CompletionsProcessor(KMContext *c) : context(c) {}

        ~CompletionsProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct ChatCompletionsProcessor : public krpc::RestfulProcessor {

        ChatCompletionsProcessor(KMContext *c) : context(c) {}

        ~ChatCompletionsProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override {

        }

        void process_with_done(const krpc::RestfulRequest *request, krpc::RestfulResponse *response, ::google::protobuf::Closure *done) override;

        bool is_wildcards() const override;

        bool need_done() const override {
            return true;
        }

        KMContext *context;
    };

    struct InfillProcessor : public krpc::RestfulProcessor {

        InfillProcessor(KMContext *c) : context(c) {}

        ~InfillProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct EmbeddingProcessor : public krpc::RestfulProcessor {

        EmbeddingProcessor(KMContext *c) : context(c) {}

        ~EmbeddingProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct RerankProcessor : public krpc::RestfulProcessor {

        RerankProcessor(KMContext *c) : context(c) {}

        ~RerankProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct TokenizeProcessor : public krpc::RestfulProcessor {

        TokenizeProcessor(KMContext *c) : context(c) {}

        ~TokenizeProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct DetokenizeProcessor : public krpc::RestfulProcessor {

        DetokenizeProcessor(KMContext *c) : context(c) {}

        ~DetokenizeProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct LoraListProcessor : public krpc::RestfulProcessor {

        LoraListProcessor(KMContext *c) : context(c) {}

        ~LoraListProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct LoraApplyProcessor : public krpc::RestfulProcessor {

        LoraApplyProcessor(KMContext *c) : context(c) {}

        ~LoraApplyProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct RootProcessor : public krpc::RestfulProcessor {

        RootProcessor() = default;

        ~RootProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    };

    struct SlotsProcessor : public krpc::RestfulProcessor {

        SlotsProcessor(KMContext *c) : context(c) {}

        ~SlotsProcessor() override = default;

        void process(const krpc::RestfulRequest *request, krpc::RestfulResponse *response) override;

        bool is_wildcards() const override;

        KMContext *context;
    private:
        void get(const krpc::RestfulRequest *request, krpc::RestfulResponse *response);
        void action(const krpc::RestfulRequest *request, krpc::RestfulResponse *response);
    };

    void setup_oai_api(KMContext *c);

    void setup_oai_ui(KMContext *c);

}  // namespace kllm
